# visual_verify_scannet.py

import os
import os.path as osp
import sys
import glob
import logging
import argparse
import numpy as np
from tqdm import tqdm

# --- Import from our visual utility file ---
from visual_util import launch_visualizer

# --- Import from the original utility file for data loading helpers ---
# Note: Assumes verify_utils.py is in the same directory or accessible in the python path.
try:
    # We assume verify_utils.py now contains your updated imread_cv2 function
    from verify_utils import setup_logging, depthmap_to_absolute_camera_coordinates, imread_cv2
except ImportError:
    print("FATAL ERROR: Could not import from 'verify_utils.py'.")
    print("Please ensure 'verify_utils.py' is in the same directory.")
    sys.exit(1)


def load_scannet_view_data(rgb_path, logger):
    """
    Loads all necessary data for a single view from the preprocessed ScanNet++ dataset format.
    """
    try:
        # Construct paths based on the ScanNet++ naming convention
        depth_path = rgb_path.replace("_rgb.jpg", "_depth.png")
        cam_path = rgb_path.replace("_rgb.jpg", "_cam.npz")

        if not all(osp.exists(p) for p in [depth_path, cam_path]):
            logger.warning(f"Missing files for {osp.basename(rgb_path)}. Skipping.")
            return None

        # Load image and depth data using the updated imread_cv2
        rgb_image = imread_cv2(rgb_path) # Assumes it returns RGB
        depthmap = imread_cv2(depth_path, -1) # Use -1 for cv2.IMREAD_UNCHANGED
        
        if rgb_image is None or depthmap is None:
            logger.warning(f"Failed to read image or depth for {rgb_path}. Skipping.")
            return None

        # --- Process data according to ScanNet++ format ---
        # 1. Handle invalid depth values
        depthmap[depthmap == 65535] = 0
        # 2. Scale depth from millimeters (uint16) to meters (float32)
        depthmap = depthmap.astype(np.float32) / 1000.0
        # 3. Normalize RGB image for the viewer
        rgb_image = rgb_image.astype(np.float32) / 255.0

        # Load camera data
        cam_file = np.load(cam_path)
        intrinsics = cam_file["intrinsics"].astype(np.float32)
        camera_pose = cam_file["pose"].astype(np.float32)

        return {
            "rgb": rgb_image, "depth": depthmap, "intrinsics": intrinsics,
            "pose": camera_pose, "path": rgb_path
        }
    except Exception as e:
        logger.error(f"Could not process view for {osp.basename(rgb_path)}. Details: {e}", exc_info=True)
        return None

# ======================================================================================
# --- Main Execution Block ---
# ======================================================================================

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Visually verify the consistency of a PREPROCESSED ScanNet++ dataset."
    )
    
    parser.add_argument(
        "--data_root", 
        type=str,
        default='/mnt/sda/scannetpp_processed',
        help="Path to the root directory of the preprocessed ScanNet++ dataset."
    )
    
    parser.add_argument(
        "--log_file", 
        type=str, 
        default="logs/visual_verification_scannet.log",
        help="Path to the output log file."
    )

    parser.add_argument(
        "--scene", 
        type=str, 
        default='8890d0a267',
        help="Optional: Specify a single scene name (e.g., '64426040a3f054a3cb452093') to visualize."
    )

    parser.add_argument(
        "--port", 
        type=int, 
        default=8080,
        help="Port to run the Viser web server on."
    )
    
    args = parser.parse_args()
    
    # Ensure the log directory exists
    log_dir = osp.dirname(args.log_file)
    if log_dir:
        os.makedirs(log_dir, exist_ok=True)
    logger = setup_logging(args.log_file)

    if not osp.isdir(args.data_root):
        logger.error(f"Provided data root '{args.data_root}' does not exist.")
        sys.exit(1)
    
    if args.scene:
        scene_dirs = [osp.join(args.data_root, args.scene)]
        if not osp.isdir(scene_dirs[0]):
            logger.error(f"Specified scene directory '{scene_dirs[0]}' does not exist.")
            sys.exit(1)
    else:
        scene_dirs = sorted([d for d in glob.glob(osp.join(args.data_root, '*')) if osp.isdir(d)])
    
    logger.info(f"Found {len(scene_dirs)} scenes to process in '{args.data_root}'.")
    
    for scene_dir in scene_dirs:
        scene_key = f"{osp.basename(scene_dir)}"
        logger.info(f"\n--- Processing scene: {scene_key} ---")

        # Find all RGB images for the scene
        rgb_paths = sorted(glob.glob(osp.join(scene_dir, "*_rgb.jpg")))[:50]
        if not rgb_paths:
            logger.warning(f"No RGB images found in {scene_dir}. Skipping.")
            continue

        # --- Prepare data for the visualizer ---
        pc_list, color_list, conf_list = [], [], []
        focals, pps, Rs, ts = [], [], [], []
        
        logger.info("Loading views and generating point clouds...")
        for rgb_path in tqdm(rgb_paths, file=sys.stdout):
            view_data = load_scannet_view_data(rgb_path, logger)
            if view_data:
                # Generate world-space point cloud for this single view
                world_points, valid_mask, _ = depthmap_to_absolute_camera_coordinates(
                    view_data["depth"], view_data["intrinsics"], view_data["pose"]
                )
                pc_list.append(world_points)
                color_list.append(view_data['rgb'])
                conf_list.append(valid_mask.astype(np.float32))

                # Aggregate camera parameters for the viewer
                intrinsics = view_data['intrinsics']
                pose = view_data['pose']
                focals.append(intrinsics[0, 0])
                pps.append((intrinsics[0, 2], intrinsics[1, 2]))
                Rs.append(pose[:3, :3])
                ts.append(pose[:3, 3])
    
        if not pc_list:
            logger.warning(f"No valid data could be loaded for {scene_key}. Skipping visualization.")
            continue

        cam_dict = {'focal': focals, 'pp': pps, 'R': Rs, 't': ts}

        # --- Launch the visualizer ---
        try:
            logger.info(f"Launching interactive visualizer for {scene_key} at http://localhost:{args.port}")
            logger.info("--> Close the browser tab or press Ctrl+C in the console to continue.")
            launch_visualizer(
                pc_list=pc_list,
                color_list=color_list,
                conf_list=conf_list,
                cam_dict=cam_dict,
                port=args.port
            )
        except KeyboardInterrupt:
            logger.info(f"\nVisualizer for {scene_key} closed, proceeding...")
        except Exception as e:
            logger.error(f"Failed to launch visualizer for {scene_key}: {e}", exc_info=True)
        break

    logger.info("\n--- Visual verification complete. ---")
